from copy import copy
from enum import Enum, auto
from itertools import count

from nanovllm.sampling_params import SamplingParams


class SequenceStatus(Enum):
    """
    WAITING：等待处理的序列
    RUNNING：正在处理的序列
    FINISHED：处理完成的序列
    """
    WAITING = auto()
    RUNNING = auto()
    FINISHED = auto()


class Sequence:
    block_size = 256 #定义每个内存块可容纳的 token 数量，用于分页存储
    counter = count() #用于生成全局唯一的序列 ID 的计数器

    def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
        self.seq_id = next(Sequence.counter)
        self.status = SequenceStatus.WAITING
        # 深拷贝输入的 token ID 列表，避免外部修改影响内部状态
        self.token_ids = copy(token_ids)
        # 记录最后一个 token，用于模型推理
        self.last_token = token_ids[-1]
        self.num_tokens = len(self.token_ids)
        self.num_prompt_tokens = len(token_ids)
        self.num_cached_tokens = 0
        # 块表，记录该序列使用的所有内存块 ID
        self.block_table = []
        self.temperature = sampling_params.temperature
        self.max_tokens = sampling_params.max_tokens
        self.ignore_eos = sampling_params.ignore_eos
    # 支持 len(sequence) 语法。
    def __len__(self):
        return self.num_tokens

    #支持通过索引访问 token_ids，允许 sequence[i] 或 sequence[start:end] 语法。
    def __getitem__(self, key):
        return self.token_ids[key]

    @property
    def is_finished(self):
        return self.status == SequenceStatus.FINISHED

    @property
    def num_completion_tokens(self): #计算已生成的补全部分的 token 数。
        return self.num_tokens - self.num_prompt_tokens

    @property
    def prompt_token_ids(self):
        return self.token_ids[:self.num_prompt_tokens]

    @property
    def completion_token_ids(self): 
        return self.token_ids[self.num_prompt_tokens:]

    @property
    def num_cached_blocks(self): #计算已缓存的块数量，用于优化内存使用。
        return self.num_cached_tokens // self.block_size

    @property
    def num_blocks(self): #计算序列需要的总块数，使用向上取整公式确保足够空间。  
        # 本身是向下取整，但通过 + self.block_size - 1 的 “补偿”，最终实现了向上取整的效果
        return (self.num_tokens + self.block_size - 1) // self.block_size

    @property
    def last_block_num_tokens(self):  #计算最后一个块中实际包含的 token 数（可能小于 block_size）
        return self.num_tokens - (self.num_blocks - 1) * self.block_size

    def block(self, i):
        assert 0 <= i < self.num_blocks
        # 返回指定块索引对应的 token 列表
        return self.token_ids[i*self.block_size: (i+1)*self.block_size]

    def append_token(self, token_id: int):
        self.token_ids.append(token_id)
        self.last_token = token_id
        self.num_tokens += 1

    def __getstate__(self):
        return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table,
                self.token_ids if self.num_completion_tokens == 0 else self.last_token)

    def __setstate__(self, state):
        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1]
        if self.num_completion_tokens == 0:
            self.token_ids = state[-1]
        else:
            self.last_token = state[-1]
